
import os

def diskUsageDictionary(path, collector=None):
    """
    returns (collector, totalbytes, subdirectories)
    collector: descendantpath --> (descendanttotal, descendantsubdirectories)
    subdirectories: dirpath --> bytecountfordirpath
    """
    total = 0
    mydirs = None
    if collector is None:
        collector = {}
    if os.path.isfile(path):
        total = os.path.getsize(path)
    elif os.path.isdir(path):
        #pr "analysing dir", path
        mydirs = {}
        filenames = os.listdir(path)
        total = 0
        filestotal = 0
        for fn in filenames:
            filepath = os.path.join(path, fn)
            (collector, subtotal, dirfiles) = diskUsageDictionary(filepath, collector)
            if dirfiles:
                mydirs[filepath] = subtotal
                #pr "subdir", (filepath, subtotal)
            else:
                filestotal += subtotal
                #pr "file", (filepath, subtotal, filestotal)
            total += subtotal
        #pr "final subtotal, filestotal", subtotal, filestotal
        mydirs["[files]"] = filestotal
    else:
        # don't know what to do with nondir/nonfile
        pass
    collector[path] = (total, mydirs)
    return (collector, total, mydirs)

if __name__=="__main__":
    target = "."
    import sys
    if len(sys.argv)>1:
        target = sys.argv[1]
    print "testing diskUsageDictionary on", repr(target)
    (c,t,d) = diskUsageDictionary(target)
    L = [ (s,n) for (n,(s,sd)) in c.items() ]
    L.sort()
    for x in L:
        print x
    print "total:", t
    if d:
        print "top level info"
        print "regular files", (target, d["[files]"])
        del d["[files]"]
        L = [ (s,p) for (p,s) in d.items() ]
        L.sort()
        for (s,p) in L:
            print "dir", (p,s)
